from Neural_Overlapping_Quantization.train_eval_model_3 import neural_network
import numpy as np
from Neural_Overlapping_Quantization.util import dir_io
import time
import os
import multiprocessing
from multiprocessing.managers import BaseManager


def train_eval_model(base, query, trainset, config):
    save_dir = '%s/Classifier_%d' % (
        config['program_train_para_dir'], config['classifier_number'])
    config['save_dir'] = save_dir
    train_model_ins = train_model_factory(config)
    # if use the learn dataset, add the learn variable in here
    train_model_ins.train(base, trainset)
    eval_result, intermediate_config = train_model_ins.eval(query)
    # train_model_ins.save()
    return eval_result, intermediate_config


def train_model_factory(config):
    _type = config['type']
    if _type == 'neural_network':
        return neural_network.NeuralNetwork(config)
    raise Exception('do not support the type of training data')
